d3d5d3
@@ -31,6 +31,7 @@
 import org.apache.hadoop.hive.conf.HiveConf;
 import org.apache.hadoop.hive.ql.CompilationOpContext;
 import org.apache.hadoop.hive.ql.ErrorMsg;
+import org.apache.hadoop.hive.ql.exec.vector.VectorizedRowBatch;
 import org.apache.hadoop.hive.ql.metadata.HiveException;
 import org.apache.hadoop.hive.ql.metadata.VirtualColumn;
 import org.apache.hadoop.hive.ql.plan.OperatorDesc;
@@ -107,9 +108,21 @@
public String getSchemaEvolutionColumnsTypes() {
    **/
   @Override
   public void process(Object row, int tag) throws HiveException {
-    if (rowLimit >= 0 && currCount++ >= rowLimit) {
-      setDone(true);
-      return;
+    if (rowLimit >= 0) {
+      if (row instanceof VectorizedRowBatch) {
+        VectorizedRowBatch batch = (VectorizedRowBatch) row;
+        if (currCount >= rowLimit) {
+          setDone(true);
+          return;
+        }
+        if (currCount + batch.size > rowLimit) {
+          batch.size = rowLimit - currCount;
+        }
+        currCount += batch.size;
+      } else if (currCount++ >= rowLimit) {
+        setDone(true);
+        return;
+      }
     }
     if (conf != null && conf.isGatherStats()) {
       gatherStats(row);
